from functools import partial

import numpy as np
import numpy.random
import torch

import sko
from scene.vrptw_opt.model import vrptw_cost
from scene.vrptw_opt.model.M_GA_TSP import GA_MTSP


class VRPTW_OPT:
    def __init__(self, args, data, vrptw_cost_func=None):
        # 导入真实数据集数据
        numpy.random.seed(args.seed)
        torch.manual_seed(args.seed)
        self.args = args
        self.data = data
        self.vrptw_cost_func = vrptw_cost_func
        self.vrptw_verify(self.args.n_dim)
        self.costfun = partial(
            vrptw_cost.vrtpw_cost,
            data=self.data,
            device=self.args.device,
        )

    def vrptw_verify(self, ndim):
        assert self.data.QT.shape[0] == ndim + 1
        assert self.data.FT.shape[0] == ndim + 1
        assert self.data.ST.shape[0] == ndim + 1
        assert self.data.QT.shape[0] == ndim + 1
        assert self.data.route_dist.shape[-2] == ndim + 1
        assert self.data.route_dist.shape[-1] == ndim + 1
        assert self.data.route_time.shape[-2] == ndim + 1
        assert self.data.route_time.shape[-1] == ndim + 1
        assert self.data.location.shape[0] == ndim + 1

    #  查找最优解
    def compute(self):
        sko.tools.set_run_mode(self.costfun, 'vectorization')
        self.ga_tsp = GA_MTSP(func=self.costfun, args=self.args)
        self.ga_tsp.to(self.args.device)
        for i in range(self.args.retry):
            self.a, self.curcost = self.ga_tsp.run()
            self.chkrouter(self.a)
            if self.curcost > 1e5:
                self.ga_tsp.crtbp()
                self.ga_tsp.to(self.args.device)
            else:
                break
        return self.curcost, self.ga_tsp

    # 检查路径的合理性
    def chkrouter(self, router, en_print=True):
        route_time = self.data.route_time
        time_window = np.array([self.data.QT, self.data.FT]).T.astype(numpy.float32)
        serve_time = np.array(self.data.ST).astype(numpy.float32)
        start_time = self.data.QT - self.data.route_time[0, 0, :]

        ans = True
        max_dely = 0
        min_early_arrive = 16384
        if en_print:
            print('check result')
        cur = 0.0
        real_wt = []
        real_tm = []
        for item in range(router.shape[0]):
            if router[item] == 0:
                cur = start_time[router[(item + 1) % router.shape[0]]]
            cur += route_time[
                np.minimum(
                    route_time.shape[0] - 1,  # 若没有找到数据则按最后一组数据（填充的平均值）计算
                    np.maximum(0,
                               cur / 30 - 1
                               ).astype(np.int16)  # 取正数部分并四舍五入，判断是哪半个小时的数据，负数取0
                ),
                router[item],  # 出发点
                router[(item + 1) % router.shape[0]]  # 到达点
            ]  # 到达时间=出发时间+路程时间
            cur = np.maximum(cur, time_window[router[(item + 1) % router.shape[0]], 0])
            real_wt.append(cur - time_window[router[(item + 1) % router.shape[0]], 0])
            wt = time_window[router[(item + 1) % router.shape[0]], 1]
            if wt >= cur:
                if en_print:
                    print('%4d->%4d: %4d pass' % (cur, wt, router[(item + 1) % router.shape[0]]), end=' ')
                min_early_arrive = min(min_early_arrive, wt - cur)
            else:
                ans = False
                max_dely = max(max_dely, - wt + cur)
                if en_print:
                    print('%4d->%4d: %4d fail' % (cur, wt, router[(item + 1) % router.shape[0]]), end=' ')
            cur += serve_time[router[(item + 1) % router.shape[0]]]  # 到达时间+服务时间=该网点出发时间
            if router[(item + 1) % router.shape[0]] == 0:
                real_tm.append(cur)
                print()

        if en_print:
            print(ans)
        return min_early_arrive, max_dely, np.mean(real_wt), np.mean(real_tm)
